package com.campus.counseling.service.impl;

import com.campus.counseling.entity.TrainingData;
import com.campus.counseling.service.DL4JModelService;
import com.campus.counseling.service.LSTMModelTrainer;
import com.campus.counseling.service.ModelEvaluationService;
import com.campus.counseling.service.TrainingDataService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.springframework.stereotype.Service;
import com.campus.counseling.model.ModelMetrics;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.stream.Collectors;

@Slf4j
@Service
@RequiredArgsConstructor
public class TrainingServiceImpl {

    private final PsyQADataLoaderImpl psyQADataLoader;
    private final DL4JModelService dl4jModelService;
    private final TrainingDataService trainingDataService;
    private final ModelEvaluationService evaluationService;

    public void trainWithPsyQA() {
        try {
            // 1. 加载训练数据
            List<TrainingData> trainingDataList = psyQADataLoader.loadPsyQAData();
            log.info("加载数据完成，总样本数: {}", trainingDataList.size());

            // 2. 划分训练集和测试集
            List<List<TrainingData>> splits = splitData(trainingDataList, 0.8);
            List<TrainingData> trainData = splits.get(0);
            List<TrainingData> testData = splits.get(1);
            
            // 3. 准备数据集
            DataSet trainSet = trainingDataService.prepareTrainingData(trainData);
            DataSet testSet = trainingDataService.prepareTrainingData(testData);
            
            // 4. 训练前评估
            log.info("训练前模型评估:");
            ModelMetrics preTrainMetrics = evaluationService.evaluateModel(testSet);
            log.info("训练前指标: {}", preTrainMetrics);

            // 5. 训练模型
            dl4jModelService.train(trainSet);
            
            // 6. 训练后评估
            log.info("训练后模型评估:");
            ModelMetrics postTrainMetrics = evaluationService.evaluateModel(testSet);
            log.info("训练后指标: {}", postTrainMetrics);
            
            // 7. 输出改进情况
            logImprovement(preTrainMetrics, postTrainMetrics);
            
        } catch (Exception e) {
            log.error("训练失败: ", e);
            throw new RuntimeException("训练失败: " + e.getMessage());
        }
    }

    private List<List<TrainingData>> splitData(List<TrainingData> data, double trainRatio) {
        Collections.shuffle(data, new Random(42));
        int trainSize = (int) (data.size() * trainRatio);
        return Arrays.asList(
            data.subList(0, trainSize),
            data.subList(trainSize, data.size())
        );
    }

    private void logImprovement(ModelMetrics pre, ModelMetrics post) {
        double mseImprovement = ((pre.getMse() - post.getMse()) / pre.getMse()) * 100;
        double accImprovement = (post.getAccuracy() - pre.getAccuracy()) * 100;
        
        log.info("模型改进情况:");
        log.info("MSE减少: {:.2f}%", mseImprovement);
        log.info("准确率提升: {:.2f}个百分点", accImprovement);
    }
}